#pragma once
#include"PreStd.h"

namespace Bonnie {
    class BestFitAllocator
    {
    public:
        using TOffset = size_t;
        using TSize = size_t;
    private:
        
        struct FreeBlockInfo;
        using FreeBlockByOffsetMap = std::map< TOffset, FreeBlockInfo>;
        using FreeBlockBySizeMap = std::map< TSize, FreeBlockByOffsetMap::iterator>;
        struct FreeBlockInfo
        {
            TSize Size;
            FreeBlockBySizeMap::iterator OrderBySizeIt;
            FreeBlockInfo(TSize size) :Size(size) {}
        };

        FreeBlockByOffsetMap mFreeBlockByOffset;
        FreeBlockBySizeMap mFreeBlockBySize;
        TSize mFreeSize;
        constexpr static  TOffset InvalidOffset = -1;
    public:
        BestFitAllocator(TSize size) :mFreeSize(size) {};

        TOffset Allocate(TSize size) {
            if (mFreeSize < size)
                return InvalidOffset;

            auto smallestBlockIt = mFreeBlockBySize.lower_bound(size);
            if (smallestBlockIt == mFreeBlockBySize.end())
                return InvalidOffset;

            auto offset = smallestBlockIt->second->first;
            auto newOffset = offset + size;
            auto newSize = smallestBlockIt->first - size;

            mFreeBlockByOffset.erase(smallestBlockIt->second);
            mFreeBlockBySize.erase(smallestBlockIt);

            if (newSize > 0)
                AddNewBlock(newOffset, newSize);

            mFreeSize -= size;
            return offset;
        }

        void Free(TOffset o, TSize s) {
            auto NextBlockIt = mFreeBlockByOffset.upper_bound(o);
            auto PrevBlockIt = NextBlockIt;
            if (PrevBlockIt != mFreeBlockByOffset.begin())
                --PrevBlockIt;
            else
                PrevBlockIt = mFreeBlockByOffset.end();

            TOffset  NewOffset;
            TSize NewSize;
            if (PrevBlockIt != mFreeBlockByOffset.end() && o == PrevBlockIt->first + PrevBlockIt->second.Size)
            {
                // PrevBlock.Offset           Offset
                // |                          |
                // |<-----PrevBlock.Size----->|<------Size-------->|
                //
                NewSize = PrevBlockIt->second.Size + s;
                NewOffset = PrevBlockIt->first;

                if (NextBlockIt != mFreeBlockByOffset.end() && o + s == NextBlockIt->first)
                {
                    // PrevBlock.Offset           Offset               NextBlock.Offset 
                    // |                          |                    |
                    // |<-----PrevBlock.Size----->|<------Size-------->|<-----NextBlock.Size----->|
                    //
                    NewSize += NextBlockIt->second.Size;
                    mFreeBlockBySize.erase(PrevBlockIt->second.OrderBySizeIt);
                    mFreeBlockBySize.erase(NextBlockIt->second.OrderBySizeIt);
                    // Delete the range of two blocks
                    ++NextBlockIt;
                    mFreeBlockByOffset.erase(PrevBlockIt, NextBlockIt);
                }
                else
                {
                    // PrevBlock.Offset           Offset                       NextBlock.Offset 
                    // |                          |                            |
                    // |<-----PrevBlock.Size----->|<------Size-------->| ~ ~ ~ |<-----NextBlock.Size----->|
                    //
                    mFreeBlockBySize.erase(PrevBlockIt->second.OrderBySizeIt);
                    mFreeBlockByOffset.erase(PrevBlockIt);
                }
            }
            else if (NextBlockIt != mFreeBlockByOffset.end() && o + s == NextBlockIt->first)
            {
                // PrevBlock.Offset                   Offset               NextBlock.Offset 
                // |                                  |                    |
                // |<-----PrevBlock.Size----->| ~ ~ ~ |<------Size-------->|<-----NextBlock.Size----->|
                //
                NewSize = s + NextBlockIt->second.Size;
                NewOffset = o;
                mFreeBlockBySize.erase(NextBlockIt->second.OrderBySizeIt);
                mFreeBlockByOffset.erase(NextBlockIt);
            }
            else
            {
                // PrevBlock.Offset                   Offset                       NextBlock.Offset 
                // |                                  |                            |
                // |<-----PrevBlock.Size----->| ~ ~ ~ |<------Size-------->| ~ ~ ~ |<-----NextBlock.Size----->|
                //
                NewSize = s;
                NewOffset = o;
            }

            AddNewBlock(NewOffset, NewSize);

            mFreeSize += s;
        }
    private:
        void AddNewBlock(TOffset o, TSize s) {
            auto newBlockIt = mFreeBlockByOffset.emplace(o, s);
            auto orderIt = mFreeBlockBySize.emplace(s, newBlockIt.first);
            newBlockIt.first->second.OrderBySizeIt = orderIt.first;
        }
    };

    class BestFitGPUAllocator :public BestFitAllocator {
    
    private:
        struct FreedAllocationInfo
        {
            TOffset Offset;
            TSize Size;
            uint64_t FrameNumber;
        };
        std::queue< FreedAllocationInfo > mStaleAllocations;
        BestFitGPUAllocator(TSize size) : BestFitAllocator(size) {};
    public:
        using SharedPtr = std::shared_ptr<BestFitGPUAllocator>;
        static SharedPtr create(TSize size) {
            return std::make_shared<BestFitGPUAllocator>(size);
        }
        void Free(TOffset o, TSize s,uint64_t frameNumber){
            // Do not release the block immediately, but add
            // it to the queue instead
            mStaleAllocations.emplace(o, s, frameNumber);
        }

        void ReleaseCompletedFrame(uint64_t numCompletedFrame) {
            // Free all allocations from the beginning of the queue that belong to completed frames
            while (!mStaleAllocations.empty() && mStaleAllocations.front().FrameNumber < numCompletedFrame)
            {
                auto& OldestAllocation = mStaleAllocations.front();
                BestFitAllocator::Free(OldestAllocation.Offset, OldestAllocation.Size);
                mStaleAllocations.pop();
            }
        }

    };
}
